import {
    BackendTypes,
    BaseAgentBackend,
    AgentBackendFactory
} from "../../../../public/agentBackend";
import { getEnv, EnvEnum } from "../../../../public/env";
import { BaseMessage } from "../../message/base";
import { ModelBackend, ModelBackendConfig } from "../base";
import { RoleMessageType } from "../../enum";
import { SendingChatMessage } from "../../message/chatMessage";
import { personConfigs } from "../../config";

interface OpenAiMessage {
    role: RoleMessageType;
    content: string;
}

export class AgentBackend extends ModelBackend {
    private key: string;
    private modelBackend: typeof BaseAgentBackend;

    constructor(config: ModelBackendConfig) {
        super(config);

        this.key = getEnv(EnvEnum.AGENT_BACKEND_API_KEY);

        this.modelBackend = AgentBackendFactory.create(config.modelMaker);
    }

    async run(messages: BaseMessage[]): Promise<BaseMessage | undefined> {
        const response = await this.request(messages);
        const assistantConfig = messages[0].getConfig();
        const content = response.message.content;

        if (!content) {
            return;
        }

        return new SendingChatMessage({
            roleName: assistantConfig.roleName,
            roleType: assistantConfig.roleType,
            roleMessageType: RoleMessageType.AgentSending,
            content: content,
            metaDict: response
        });
    }

    getTokenLimit(): number {
        return this.modelBackend.getTokenLimit(this.modelType);
    }

    private async request(
        messages: BaseMessage[],
        retryCount: number = 0
    ): Promise<BackendTypes.ChatCompletionResponse> {
        try {
            const response = await this.modelBackend.getChatCompletion(
                {
                    apiKey: this.key,
                    model: this.modelType,
                    messages: getMessages(messages),
                    ...this.modelConfig
                },
                this.modelConfig
            );
            const content = response.message.content;
            if (!content) {
                throw new Error("content为空");
            }
            return response;
        } catch (error) {
            if (retryCount < personConfigs.retryNum) {
                // 随机延时重试 1s - 2s
                const randomTime = Math.floor(Math.random() * 1000) + 1000;
                await new Promise((resolve) =>
                    setTimeout(() => {
                        resolve(null);
                    }, randomTime)
                );
                return await this.request(messages, retryCount + 1);
            }
            throw Error(
                `无法获取 chat completion 结果，请检查网络连接或重试; ${(error as Error).message}`
            );
        }
    }
}

function getMessages(messages: BaseMessage[]): OpenAiMessage[] {
    return messages.map((message) => {
        return {
            role: message.roleMessageType,
            content: message.content
        };
    });
}
